from typing import List, Iterator, Dict, Any

class Optimizer:
    """
    优化器基类
    """
    def __init__(self, params, defaults):
        """
        初始化优化器
        
        Args:
            params: 要优化的参数
            defaults: 默认超参数
        """
        self.defaults = defaults
        self.state = dict()  # 优化器状态
        self.param_groups = []
        
        param_groups = list(params)
        if len(param_groups) == 0:
            raise ValueError("优化器没有参数")
            
        if not isinstance(param_groups[0], dict):
            param_groups = [{'params': param_groups}]
            
        for param_group in param_groups:
            self.add_param_group(param_group)
    
    def __repr__(self):
        format_string = self.__class__.__name__ + ' ('
        for i, group in enumerate(self.param_groups):
            format_string += '\n'
            format_string += f'Parameter Group {i}\n'
            for key in sorted(group.keys()):
                if key != 'params':
                    format_string += f'    {key}: {group[key]}\n'
        format_string += ')'
        return format_string
    
    def state_dict(self):
        """
        返回优化器状态字典
        """
        return {
            'state': self.state,
            'param_groups': self.param_groups,
        }
    
    def load_state_dict(self, state_dict):
        """
        加载状态字典
        
        Args:
            state_dict: 优化器状态字典
        """
        self.state = state_dict['state']
        self.param_groups = state_dict['param_groups']
    
    def zero_grad(self):
        """
        清除所有参数的梯度
        """
        for group in self.param_groups:
            for p in group['params']:
                # 检查参数是否有grad属性
                if hasattr(p, 'grad') and p.grad is not None:
                    p.grad = None
    
    def step(self):
        """
        执行单步参数更新，需要在子类中实现
        """
        raise NotImplementedError
    
    def add_param_group(self, param_group):
        """
        添加参数组
        
        Args:
            param_group: 包含参数和超参数的字典
        """
        assert isinstance(param_group, dict), "param group必须是字典"
        
        params = param_group['params']
        if isinstance(params, Iterator):
            params = list(params)
            
        param_group['params'] = params
        
        for name, default in self.defaults.items():
            if name not in param_group:
                param_group[name] = default
                
        self.param_groups.append(param_group) 